{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Technical Indicators Demo\n",
    "\n",
    "## Download example datasets\n",
    "\n",
    "Before getting started, let's download the example datasets if not present."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset is already present. No need to re-download it.\n"
     ]
    }
   ],
   "source": [
    "! ((test ! -f '../data/stock_price_hist.csv.gz' ||  test ! -f '../data/security_master.csv.gz') && \\\n",
    "  cd ../.. && bash download_data.sh) || echo \"Dataset is already present. No need to re-download it.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bqplot import OHLC, LinearScale, DateScale, Axis, Figure, Bars, Lines, Tooltip\n",
    "from bqplot.colorschemes import CATEGORY20\n",
    "import datetime\n",
    "import ipywidgets as widgets\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys; sys.path.append('../..')\n",
    "\n",
    "from gquant.dataframe_flow import TaskSpecSchema\n",
    "\n",
    "\n",
    "task_load_csv_data = {\n",
    "    TaskSpecSchema.task_id: \"load_csv_data\",\n",
    "    TaskSpecSchema.node_type: \"CsvStockLoader\",\n",
    "    TaskSpecSchema.conf: {\"path\": \"../data/stock_price_hist.csv.gz\"},\n",
    "    TaskSpecSchema.inputs: []\n",
    "}\n",
    "\n",
    "task_sort = {\n",
    "    TaskSpecSchema.task_id: \"sort\",\n",
    "    TaskSpecSchema.node_type: \"SortNode\",\n",
    "    TaskSpecSchema.conf: {\"keys\": ['asset', 'datetime']},\n",
    "    TaskSpecSchema.inputs: [\"load_csv_data\"]\n",
    "}\n",
    "\n",
    "task_stock_symbol = {\n",
    "    TaskSpecSchema.task_id: \"stock_symbol\",\n",
    "    TaskSpecSchema.node_type: \"StockNameLoader\",\n",
    "    TaskSpecSchema.conf: {\"path\": \"../data/security_master.csv.gz\"},\n",
    "    TaskSpecSchema.inputs: []\n",
    "}    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import warnings; warnings.simplefilter(\"ignore\")\n",
    "\n",
    "from gquant.dataframe_flow import TaskGraph\n",
    "\n",
    "task_list = [task_load_csv_data, task_sort]\n",
    "task_graph = TaskGraph(task_list)\n",
    "\n",
    "action = \"load\" if os.path.isfile('./.cache/load_csv_data.hdf5') else \"save\"\n",
    "\n",
    "df = task_graph.run(outputs=['sort'], replace={'load_csv_data': {action: True}})[0]\n",
    "\n",
    "def one_stock(df, stock_id):\n",
    "    return df.query('asset==%s' % stock_id)\n",
    "\n",
    "def slice_stock(df, year):\n",
    "    beg_date = datetime.datetime.strptime(str(year)+'-01-01', '%Y-%m-%d')\n",
    "    end_date = datetime.datetime.strptime(str(int(year)+1)+'-01-01', '%Y-%m-%d')\n",
    "    return df.query('datetime<@end_date and datetime>=@beg_date')\n",
    "\n",
    "indicator_lists = ['Accumulation Distribution', 'ADMI', 'Average True Range', 'Bollinger Bands',\n",
    "                   'Chaikin Oscillator', 'Commodity Channel Index', 'Coppock Curve', 'Donchian Channel',\n",
    "                   'Ease of Movement', 'EWA', 'Force Index', 'Keltner Channel', 'KST Oscillator', 'MA', 'MACD',\n",
    "                   'Mass Index', 'Momentum', 'Money Flow Index', 'On Balance Volume', 'Parabolic SAR',\n",
    "                   'Rate of Change', 'RSI', 'Stochastic Oscillator D', 'Stochastic Oscillator K', 'TRIX',\n",
    "                   'True Strength Index', 'Ultimate Oscillator', 'Vortex Indicator',]\n",
    "\n",
    "task_stocks_list = [task_stock_symbol]\n",
    "task_stocks_graph = TaskGraph(task_stocks_list)\n",
    "list_stocks = task_stocks_graph.run(outputs=['stock_symbol'])[0].to_pandas().set_index('asset_name').to_dict()['asset']\n",
    "\n",
    "main_figure_height='300px'\n",
    "indicator_figure_height='200px'\n",
    "figure_width = '1000px'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f3a3621fc2104f31b96493bbe902121c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HBox(children=(Dropdown(description='Add stock', options=('MIDAX', 'MMELX', 'BBDO', 'BSAC', 'BI…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "add_stock_selector = widgets.Dropdown(options=list_stocks.keys(), value=None, description=\"Add stock\")\n",
    "year_selector = widgets.IntSlider(description=\"All Year\", continuous_update=False)\n",
    "year_selectors = []\n",
    "\n",
    "def get_figure(selected, df):\n",
    "    this_stock = one_stock(df, list_stocks[selected])\n",
    "    this_stock_store = [this_stock]\n",
    "    stock_selector = widgets.Dropdown(options=list_stocks.keys(), value=add_stock_selector.value, description=\"stock\")\n",
    "    indicator_selector = widgets.Dropdown(options=indicator_lists, value=None, description=\"Indicator\")\n",
    "    min_year = this_stock.datetime.dt.year.min()\n",
    "    max_year = this_stock.datetime.dt.year.max()\n",
    "    year_selector = widgets.IntSlider(min=min_year, max=max_year, description=\"Year\", continuous_update=False)\n",
    "    year = year_selector.value\n",
    "    stock = slice_stock(this_stock, year)\n",
    "    sc = LinearScale()\n",
    "    sc2 = LinearScale()\n",
    "    dt_scale = DateScale()\n",
    "    ax_x = Axis(label='Date', scale=dt_scale)\n",
    "    ax_y = Axis(label='Price', scale=sc, orientation='vertical', tick_format='0.0f')\n",
    "    # Construct the marks\n",
    "    ohlc = OHLC(x=stock.datetime, y=stock[['open','high','low', 'close']].as_gpu_matrix(), marker='candle', scales={'x': dt_scale, 'y': sc}, format='ohlc',\n",
    "                stroke='blue', display_legend=True, labels=[selected])\n",
    "    bar = Bars(x=stock.datetime, y=stock.volume, \n",
    "           scales={'x': dt_scale, 'y': sc2}, padding=0.2)\n",
    "    def_tt = Tooltip(fields=['x', 'y'], formats=['%Y-%m-%d', '.2f'])\n",
    "    bar.tooltip = def_tt\n",
    "    bar.interactions = {\n",
    "        'legend_hover': 'highlight_axes',\n",
    "        'hover': 'tooltip', \n",
    "        'click': 'select',\n",
    "    }\n",
    "    sc.min = stock.close.min() - 0.3 * (stock.close.max() - stock.close.min()) \n",
    "    sc.max = stock.close.max()\n",
    "    sc2.max = stock.volume.max()*4.0\n",
    "    dt_scale.min = pd.Timestamp('%d-1-1' % year)\n",
    "    dt_scale.max = pd.Timestamp('%d-1-1' % (year + 1))\n",
    "    f = Figure(axes=[ax_x, ax_y], marks=[ohlc, bar], fig_margin={\"top\":0, \"bottom\":60, \"left\":60, \"right\":60})\n",
    "    f.layout.height = main_figure_height\n",
    "    f.layout.width = figure_width\n",
    "    para_selectors = widgets.VBox([])\n",
    "    color_id = [0]\n",
    "    \n",
    "    def update_graph(stock):\n",
    "        with bar.hold_trait_notifications() as bc, ohlc.hold_trait_notifications() as oc:\n",
    "            ohlc.y = stock[['open','high','low', 'close']].as_gpu_matrix()\n",
    "            ohlc.x = stock.datetime\n",
    "        \n",
    "            bar.y = stock.volume\n",
    "            bar.x = stock.datetime\n",
    "    \n",
    "            sc.min = stock.close.min() - 0.3 * (stock.close.max() - stock.close.min()) \n",
    "            sc.max = stock.close.max()\n",
    "            sc2.max = stock.volume.max()*4.0\n",
    "            dt_scale.min = pd.Timestamp('%d-1-1' % year_selector.value)\n",
    "            dt_scale.max = pd.Timestamp('%d-1-1' % (year_selector.value + 1))\n",
    "            update_range()\n",
    "    \n",
    "    def year_selection(*stock):\n",
    "        stock = slice_stock(this_stock_store[0], year_selector.value)\n",
    "        update_graph(stock)\n",
    "    \n",
    "    def stock_selection(*stock):\n",
    "        this_stock_store[0] = one_stock(df, list_stocks[stock_selector.value])\n",
    "        year_selector.min = this_stock_store[0].Dte.to_array().min().astype(datetime.datetime).year\n",
    "        year_selector.max = this_stock_store[0].Dte.to_array().max().astype(datetime.datetime).year\n",
    "        stock = slice_stock(this_stock_store[0], year_selector.value)\n",
    "        ohlc.labels = [stock_selector.value]\n",
    "        update_graph(stock)\n",
    "        \n",
    "    def update_figure_(stock, objects):\n",
    "        line = objects[0]\n",
    "        with line.hold_trait_notifications():\n",
    "            line.y = stock['out']\n",
    "            line.x = stock.datetime\n",
    "            \n",
    "    def add_new_indicator(new_fig):\n",
    "        # add new figure\n",
    "        # take the axis from the fig\n",
    "        empty = {\"top\":0, \"bottom\":0, \"left\":60, \"right\":60}\n",
    "        axis_margin = {\"top\":0, \"bottom\":60, \"left\":60, \"right\":60}\n",
    "        axes_copy = multiple_figs.children[-1].axes.copy() \n",
    "        multiple_figs.children[-1].fig_margin = empty\n",
    "        to_be_removed = axes_copy[0]\n",
    "        axes_copy.remove(to_be_removed)\n",
    "        multiple_figs.children[-1].axes = axes_copy\n",
    "        new_axes = new_fig.axes.copy()\n",
    "        new_fig.axes = [to_be_removed] + new_axes\n",
    "        new_fig.fig_margin = axis_margin\n",
    "        # add new figure\n",
    "        multiple_figs.children += (new_fig,)\n",
    "        \n",
    "    def indicator_selection(*stock):\n",
    "        if indicator_selector.value is None:\n",
    "            return\n",
    "        \n",
    "        color_id[0] = (color_id[0] + 1) % len(CATEGORY20)\n",
    "        \n",
    "        def setup_indicator(get_para_widgets, get_parameters, process_outputs, create_figure, update_figure, indicator_fun):\n",
    "            with out:                    \n",
    "                def update_df(para_selector_widgets):\n",
    "                    my_stock = this_stock_store[0]\n",
    "                    parameters = get_parameters(my_stock, para_selector_widgets)\n",
    "                    output = indicator_fun(*parameters)\n",
    "                    stock_df = process_outputs(output, my_stock)\n",
    "                    stock = slice_stock(stock_df, year_selector.value)\n",
    "                    return stock\n",
    "                para_selector_widgets = get_para_widgets()   \n",
    "                para_selectors.children += tuple(para_selector_widgets)                    \n",
    "                stock = update_df(para_selector_widgets)\n",
    "                figs = create_figure(stock=stock, dt_scale=dt_scale, sc=sc, color_id=color_id, f=f, \n",
    "                                     indicator_figure_height=indicator_figure_height, figure_width=figure_width,\n",
    "                                     add_new_indicator=add_new_indicator)\n",
    "\n",
    "                def update_para(*para):\n",
    "                    stock = update_df(para_selector_widgets)\n",
    "                    update_figure(stock, figs)\n",
    "                    \n",
    "                for selector in para_selector_widgets:\n",
    "                    selector.observe(update_para, 'value')\n",
    "                year_selector.observe(update_para, 'value')\n",
    "                stock_selector.observe(update_para, 'value')\n",
    "\n",
    "        with out:\n",
    "            update_figure = update_figure_\n",
    "            if indicator_selector.value=='Accumulation Distribution':\n",
    "                from viz.accumulation_distribution import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun              \n",
    "\n",
    "            elif indicator_selector.value=='ADMI':\n",
    "                from viz.admi import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Average True Range':\n",
    "                from viz.average_true_range import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Bollinger Bands':\n",
    "                from viz.bollinger_bands import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun, update_figure\n",
    "\n",
    "            elif indicator_selector.value=='Chaikin Oscillator':\n",
    "                from viz.ch_oscillator import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Commodity Channel Index':\n",
    "                from viz.commodity_channel_index import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Coppock Curve':\n",
    "                from viz.coppock_curve import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun    \n",
    "\n",
    "            elif indicator_selector.value=='Donchian Channel':\n",
    "                from viz.donchian_channel import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun      \n",
    "\n",
    "            elif indicator_selector.value=='Ease of Movement':\n",
    "                from viz.ease_of_movement import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun             \n",
    "\n",
    "            elif indicator_selector.value=='EWA':\n",
    "                from viz.ewa import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Force Index':\n",
    "                from viz.force_index import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Keltner Channel':\n",
    "                from viz.keltner_channel import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun       \n",
    "\n",
    "            elif indicator_selector.value=='KST Oscillator':\n",
    "                from viz.kst_oscillator import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='MA':\n",
    "                from viz.ma import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='MACD':           \n",
    "                from viz.macd import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun, update_figure\n",
    "\n",
    "            elif indicator_selector.value=='Mass Index':\n",
    "                from viz.mass_index import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "            elif indicator_selector.value=='Momentum':\n",
    "                from viz.momentum import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Money Flow Index':\n",
    "                from viz.money_flow_index import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='On Balance volume':\n",
    "                from viz.on_balance_volume import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Parabolic SAR':\n",
    "                from viz.parabolic_sar import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun, update_figure\n",
    "\n",
    "            elif indicator_selector.value=='Rate of Change':\n",
    "                from viz.rate_of_change import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='RSI':\n",
    "                from viz.rsi import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Stochastic Oscillator D':\n",
    "                from viz.stochastic_oscillator_d import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='Stochastic Oscillator K':\n",
    "                from viz.stochastic_oscillator_k import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun            \n",
    "\n",
    "            elif indicator_selector.value=='TRIX':\n",
    "                from viz.trix import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "            elif indicator_selector.value=='True Strength Index':\n",
    "                from viz.true_strength_index import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun              \n",
    "\n",
    "            elif indicator_selector.value=='Ultimate Oscillator':\n",
    "                from viz.ultimate_oscillator import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun                  \n",
    "                \n",
    "            elif indicator_selector.value=='Vortex Indicator':\n",
    "                from viz.vortex_indicator import get_para_widgets, get_parameters, process_outputs, create_figure, indicator_fun\n",
    "\n",
    "        setup_indicator(get_para_widgets, get_parameters, process_outputs, create_figure, update_figure, indicator_fun)\n",
    "\n",
    "        indicator_selector.value=None\n",
    "\n",
    "    year_selector.observe(year_selection, 'value')\n",
    "    stock_selector.observe(stock_selection, 'value')\n",
    "    indicator_selector.observe(indicator_selection, 'value')\n",
    "    multiple_figs = widgets.VBox([f])\n",
    "    return multiple_figs, year_selector, stock_selector, indicator_selector, para_selectors\n",
    "\n",
    "def stock_selection(*stock):\n",
    "    if add_stock_selector.value is None:\n",
    "        return\n",
    "    f, year_selector, stock_selector, indicator_selector, para_selectors = get_figure(add_stock_selector.value, df)\n",
    "    vbox = w.children[1]\n",
    "    vbox.children += (widgets.HBox([f, widgets.VBox([year_selector, stock_selector, indicator_selector, para_selectors])]),) \n",
    "    year_selectors.append(year_selector)\n",
    "    update_range()\n",
    "    add_stock_selector.value = None\n",
    "\n",
    "def update_range():\n",
    "    min_vals = []\n",
    "    max_vals = []\n",
    "    for i in year_selectors:\n",
    "        min_vals.append(i.min)\n",
    "        max_vals.append(i.max)\n",
    "    minV = max(min_vals)\n",
    "    maxV = min(max_vals)\n",
    "    if minV<=maxV:\n",
    "        year_selector.disabled = False\n",
    "        year_selector.max = maxV   \n",
    "        year_selector.min = minV\n",
    "    else:\n",
    "        year_selector.disabled = True\n",
    "    \n",
    "\n",
    "out = widgets.Output(layout={'border': '1px solid black'})\n",
    "\n",
    "def update_all_ranges(*arg):\n",
    "    for i in year_selectors:\n",
    "        i.value = year_selector.value\n",
    "    \n",
    "add_stock_selector.observe(stock_selection, 'value')\n",
    "year_selector.observe(update_all_ranges, 'value')\n",
    "selectors = widgets.HBox([add_stock_selector, year_selector])\n",
    "w = widgets.VBox([selectors, widgets.VBox([])])\n",
    "w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd37ad79dadd4362a1266fda125ee788",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output(layout=Layout(border='1px solid black'))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
